package com.niodata.md.rdb.crawler;

import com.niodata.md.rdb.model.DataBase;
import com.niodata.md.rdb.model.MetaType;
import com.niodata.md.rdb.model.Table;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.List;
import org.apache.commons.lang3.ArrayUtils;

/**
 * table retrieval, retrieve all table names from the parent path.
 */
public class TableRetrieval extends AbstractRetrieval {

  private TableRetrievalOptions options = TableRetrievalOptions.retrievalWithDataAndColumnInfo();

  @Override
  public List<MetaDataRetrieval> retrieveMetaData() throws Exception {
    String schemaPattern = this.options.getSchemaPattern();
    if (schemaPattern == null || schemaPattern.equals("")) {
      schemaPattern = "%";
    }
    ResultSet rs = this.getMetaData().getTables(
        this.getCatalog().getName(),
        schemaPattern,
        "%",
        new String[] {"TABLE", "VIEW"});
    List<Table> tables = new ArrayList<>();
    List<MetaDataRetrieval> subMetaDataRetrievals = new ArrayList<>();
    while (rs.next()) {
      Table table = new Table();
      table.setCatalog(this.getCatalog().getName());
      table.setDatabase(this.getCatalog().getName());
      table.setName(rs.getString("TABLE_NAME"));
      table.setType(rs.getString("TABLE_TYPE"));
      if (!table.getName().matches(options.getNamePattern())) {
        continue;
      }
      if (ArrayUtils.contains(options.getNameExclusions(), table.getName())) {
        continue;
      }
      table.setSchema(rs.getString("TABLE_SCHEM"));
      table.setRemarks(rs.getString("REMARKS"));
      tables.add(table);
      String[] path = new String[]{
          this.getCatalog().getName(),
          this.getCatalog().getName(),
          table.getName()};


      if (options.isWithColumnInfo()) {
        ColumnRetrieval columnRetrieval = new ColumnRetrieval();
        columnRetrieval.copyContext(this);
        columnRetrieval.setParentPath(path);
        subMetaDataRetrievals.add(columnRetrieval);
      }

      if (options.isWithDataInfo()) {
        AbstractRetrieval rowCountRetrieval = new RowCountRetrieval();
        Class<? extends AbstractRetrieval> clazz =
            MetadataRetrievalRegistry.getRetrievalClass(this.getDataSourceInfo().getDbCategory(), MetaType.ROW_COUNT);
        if (clazz != null) {
          rowCountRetrieval = clazz.newInstance();
        }
        rowCountRetrieval.copyContext(this);
        rowCountRetrieval.setParentPath(path);
        subMetaDataRetrievals.add(rowCountRetrieval);
      }

      IndexRetrieval indexRetrieval = new IndexRetrieval();
      indexRetrieval.copyContext(this);
      indexRetrieval.setParentPath(path);
      subMetaDataRetrievals.add(indexRetrieval);

      CreateSqlRetrieval createSqlRetrieval = new CreateSqlRetrieval();
      createSqlRetrieval.copyContext(this);
      createSqlRetrieval.setParentPath(path);
      subMetaDataRetrievals.add(createSqlRetrieval);

      if (options.isWithDataInfo()) {
        AbstractRetrieval sizeRetrieval = new TableStorageSizeRetrieval();
        Class<? extends AbstractRetrieval> sizeRetrievalClass =
            MetadataRetrievalRegistry.getRetrievalClass(this.getDataSourceInfo().getDbCategory(), MetaType.TABLE_SIZE);
        if (sizeRetrievalClass != null) {
          sizeRetrieval = sizeRetrievalClass.newInstance();
        }
        sizeRetrieval.copyContext(this);
        sizeRetrieval.setParentPath(path);
        subMetaDataRetrievals.add(sizeRetrieval);
      }
    }
    DataBase dataBase = super.getCatalog().getDataBase(this.getParentPath());
    dataBase.setTableCount(tables.size());
    dataBase.setTables(tables);

    return subMetaDataRetrievals;
  }

  /**
   * set table retrieval options.
   *
   * @param options table retrieval options
   */
  public void withTableRetrievalOptions(TableRetrievalOptions options) {
    this.options = options;
  }
}
